package tree

// TreeNode 树节点接口
type TreeNode[T any] interface {
	GetID() uint              // 获取节点ID
	GetParentID() uint        // 获取父节点ID
	SetChildren(children []T) // 设置子节点
}

// Tree 工具类
type Tree[T TreeNode[T]] struct {
	Nodes []T
}

// NewTree 创建树结构实例
func NewTree[T TreeNode[T]](nodes []T) *Tree[T] {
	return &Tree[T]{Nodes: nodes}
}

// BuildTree 构建树形结构，parentId 表示根节点的父ID，通常是 0 或其他初始值
func (t *Tree[T]) BuildTree(parentId uint) []T {
	// 创建一个节点ID到节点指针的映射
	nodeMap := make(map[uint][]T)

	// 将节点按 ParentID 分组
	for i := range t.Nodes {
		node := t.Nodes[i]
		pid := node.GetParentID()
		nodeMap[pid] = append(nodeMap[pid], node)
	}

	// 递归构建树
	var build func(parentID uint) []T
	build = func(parentID uint) []T {
		children := nodeMap[parentID]
		result := make([]T, len(children))

		for i, child := range children {
			// 复制节点到结果中
			result[i] = child
			// 设置子节点
			childrenNodes := build(child.GetID())
			result[i].SetChildren(childrenNodes)
		}

		return result
	}

	// 构建树并返回
	return build(parentId)
}
